# -*- coding: utf-8 -*-
"""
Created on 2018 3.30
@author: hugh
"""

import numpy as np
import os
import config
import pickle
import dlib
from skimage import io
from sklearn.model_selection import train_test_split


class LoadImage(object):
	
	def __init__(self, img_dir, train_label_name_data=config.train_label_name_data,
					train_face_descriptor_data=config.train_face_descriptor_data):
		"""
		加载训练图片数据，建立标签至人名字典数据和标签至128维人脸特征点的list数据
		img_dir：训练图片数据目录，适用于二级目录   img_dir/图片类别文件名/图片（.jpg）
		train_name_label_data： 标签至人名的字典数据文件
		train_face_descriptor_data：标签至128维人脸特征点的字典文件		
		"""
		if not os.path.exists(train_label_name_data) or not os.path.exists(train_face_descriptor_data):
			self._detect_save_face_descriptor(img_dir, train_label_name_data, train_face_descriptor_data)
		else:
			self.train_label_name_dict = read_obj_from_file(train_label_name_data)
			self.train_data_list = read_obj_from_file(train_face_descriptor_data)
		self.num_train_face = len(self.train_data_list)


	def get_train_test_data(self, test_size=0.05, random_state=26):
		"""获取训练数据和测试数据"""
		train_x = []
		train_y = []
		for (label, train_data) in self.train_data_list:
			train_x.append(train_data)
			train_y.append(label)
		train_x = np.array(train_x)
		train_y = np.array(train_y)
		#划分数据集为训练集合测试集
		X_train, X_test, y_train, y_test = train_test_split(train_x, train_y, test_size=test_size,
																			random_state=random_state)
		return X_train, X_test, y_train, y_test

	def get_feture_vector_from_image(self, img_path):
		"""从单张图片中获得检测到的人脸（-1，128）维特征向量"""
		# 加载人脸检测的模型
		detector = dlib.get_frontal_face_detector()
		# 加载人脸特征点预测模型
		sp = dlib.shape_predictor(config.predictor_path)
		# 加载人脸识别模型
		facerec = dlib.face_recognition_model_v1(config.face_rec_model_path)
		# 从文件读取图片
		face_descriptors = []
		img = io.imread(img_path)
		# 使用detector进行人脸检测 dets为返回的结果
		dets = detector(img, 1)
		for d in dets:
			shape = sp(img, d)
			face_descriptor = facerec.compute_face_descriptor(img, shape)
			face_descriptors.append(face_descriptor)
		return np.array(face_descriptors), dets

	def _detect_save_face_descriptor(self, img_dir, train_label_name_data, train_face_descriptor_data,
										predictor_path=config.predictor_path,
										face_rec_model_path=config.face_rec_model_path):
		"""检测识别人脸，建立并保存标签至人名字典数据和标签至128维人脸特征点的list数据至文件"""
		self.train_data_list = []
		self.train_label_name_dict = {}
		# 加载人脸检测的模型
		detector = dlib.get_frontal_face_detector()
		#加载人脸特征点预测模型
		sp = dlib.shape_predictor(predictor_path)
		#加载人脸识别模型
		facerec = dlib.face_recognition_model_v1(face_rec_model_path)
		
		#建立训练标签和人名的对应关系字典，并保存至数据文件中
		img_all_path = os.listdir(img_dir)
		for train_label, train_foldname in enumerate(img_all_path):
			self.train_label_name_dict.setdefault(train_label,train_foldname)
			print("train_foldname:{},train_label:{}".format(train_foldname, train_label))
		save_obj_to_file(self.train_label_name_dict, train_label_name_data)
		# 建立人名和标签的字典，方便提取根据人名找到对应的标签
		self.train_name_label_dict = {name: label for label, name in self.train_label_name_dict.items()}
		index = 1
		for (path, dirnames, filenames) in os.walk(img_dir):
			if len(filenames) != 0:
				train_foldname = path[path.rfind('/')+1:]
				train_label = self.train_name_label_dict.get(train_foldname)
			for filename in filenames:
				if filename.endswith('.jpg'):
					img_path = path + '/' + filename
					# 从文件读取图片
					img = io.imread(img_path)
					
					# 使用detector进行人脸检测 dets为返回的结果
					dets = detector(img, 1)
					num_faces = len(dets)
					# 去除干扰的图片（当一张训练图中检测到有多于一个人脸）
					if num_faces > 1 or num_faces == 0:
						continue
					
					for d in dets:
						# 在方框d中获取脸部的标志
						shape = sp(img, d)					
					# 计算出128维人脸特征向量。
					# 一般而言，如果两张脸图的特征向量的欧式距离小于0.6，
					# 则认为他们是同一个人，否则，他们就是不同的人。
					face_descriptor = facerec.compute_face_descriptor(img, shape)
					#将标签和人脸特征向量保存到list
					self.train_data_list.append((train_label,face_descriptor))
					print('Being processed picture %s' % index)
					index += 1
		#保存标签与人脸特征点的list
		save_obj_to_file(self.train_data_list, train_face_descriptor_data)

def save_obj_to_file(obj, filename):
	"""保存数据至指定文件"""
	with open(filename, 'wb') as f:
		pickle.dump(obj,f)

def read_obj_from_file(filename):
	"""从指定文件读取数据"""
	with open(filename, 'rb') as f:
		obj = pickle.load(f)
	return obj
